import logging
import threading
import time
import zmq
from zhelpers import zpipe
from kvsimple import KVMsg

# 请求超时时间
GLOBAL_TIMEOUT = 4000
# 判定服务死亡的时间
SERVER_TTL = 5.0
# 服务器个数
SERVER_MAX = 2

logging.basicConfig(format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)


class Clone:
    ctx = None
    pipe = None
    agent = None
    _subtree = None

    def __init__(self):
        self.ctx = zmq.Context()
        self.pipe, peer = zpipe(self.ctx)
        self.agent = threading.Thread(target=clone_agent, args=(self.ctx, peer))
        self.agent.daemon = True
        self.agent.start()

    @property
    def subtree(self):
        return self._subtree

    @subtree.setter
    def subtree(self, subtree):
        self._subtree = subtree
        logging.info("client Sends [SUBTREE][%s] to the agent" % subtree)
        self.pipe.send_multipart([b"SUBTREE", subtree])

    def connect(self, address, port):
        logging.info("client Sends [CONNECT][%s][%s] to the agent" % (address, port))
        self.pipe.send_multipart([b"CONNECT", (address.encode('utf-8')) if isinstance(address, str) else address, b"%d" % port])

    def set(self, key, value, ttl=0):
        # logging.info("client Sends [SET][%s][%s][%d] to the agent" % (key, value, ttl))
        self.pipe.send_multipart([b"SET", key, value, b'%i' % ttl])

    def get(self, key):
        # logging.info("client Sends [GET][%s] to the agent and waits for a value response" % key)
        self.pipe.send_multipart([b"GET", key])
        try:
            reply = self.pipe.recv_multipart()
        except KeyboardInterrupt:
            return
        else:
            return reply[0]


class CloneServer:
    address = None
    port = None
    snapshot = None
    subscriber = None
    expirt = 0
    requests = 0

    def __init__(self, ctx, address, port, subtree):
        self.address = address
        self.port = port

        self.snapshot = ctx.socket(zmq.DEALER)
        self.snapshot.linger = 0
        self.snapshot.connect("%s:%i" % (address.decode('utf-8'), port))

        self.subscriber = ctx.socket(zmq.SUB)
        self.subscriber.linger = 0
        self.subscriber.setsockopt(zmq.SUBSCRIBE, subtree)
        self.subscriber.setsockopt(zmq.SUBSCRIBE, b'HUGZ')
        self.subscriber.connect("%s:%i" % (address.decode('utf-8'), port+1))


STATE_INITIAL = 0
STATE_SYNCING = 1
STATE_ACTIVE = 2


class CloneAgent:
    ctx = None
    pipe = None
    subtree = ''
    servers = None
    state = 0
    cur_server = 0
    sequence = 0
    kvmap = None
    publisher = None

    def __init__(self, ctx, pipe):
        self.ctx = ctx
        self.pipe = pipe
        self.subtree = ''
        self.servers = []
        self.state = STATE_INITIAL
        self.kvmap = {}
        self.publisher = ctx.socket(zmq.PUB)

    def control_message(self):
        msg = self.pipe.recv_multipart()
        command = msg.pop(0)

        if command == b"SUBTREE":
            self.subtree = msg[0]
        elif command == b"CONNECT":
            address = msg.pop(0)
            port = int(msg.pop(0))
            if len(self.servers) < SERVER_MAX:
                self.servers.append(CloneServer(self.ctx, address, port, self.subtree))
                self.publisher.connect("%s:%i" % (address.decode('utf-8'), port+2))
            else:
                logging.error("E: too many servers (max. %i)" % SERVER_MAX)
        elif command == b"SET":
            key, value, sttl = msg
            ttl = int(sttl)

            kvmsg = KVMsg(0, key=key.decode('utf-8'), body=value.decode('utf-8'))
            kvmsg.store(self.kvmap)
            if ttl:
                kvmsg[b'ttl'] = sttl
            kvmsg.send(self.publisher)
        elif command == b"GET":
            key = msg[0]
            value = self.kvmap.get(key)
            self.pipe.send(value.body if value else '')


def clone_agent(ctx, pipe):
    agent = CloneAgent(ctx, pipe)
    server = None

    while True:
        poller = zmq.Poller()
        poller.register(agent.pipe, zmq.POLLIN)
        poll_timer = None
        server_socket = None

        if agent.state == STATE_INITIAL:
            if agent.servers:
                server = agent.servers[agent.cur_server]
                logging.info("I: waiting for server at %s:%d..." % (server.address, server.port))
                if server.requests < 2:
                    server.snapshot.send_multipart([b"ICANHAZ?", agent.subtree])
                    server.requests += 1
                server.expiry = time.time() + SERVER_TTL
                agent.state = STATE_SYNCING
                server_socket = server.snapshot
        elif agent.state == STATE_SYNCING:
            server_socket = server.snapshot
        elif agent.state == STATE_ACTIVE:
            server_socket = server.subscriber

        if server_socket:
            poller.register(server_socket, zmq.POLLIN)

        if server is not None:
            poll_timer = 1e3 * max(0, server.expiry - time.time())

        try:
            items = dict(poller.poll(poll_timer))
        except KeyboardInterrupt:
            raise
            break

        if agent.pipe in items:
            agent.control_message()
        elif server_socket in items:
            kvmsg = KVMsg.recv(server_socket)

            server.expiry = time.time() + SERVER_TTL
            if agent.state == STATE_SYNCING:
                server.requests = 0
                if kvmsg.key == "KTHXBAI":
                    agent.sequence = kvmsg.sequence
                    agent.state = STATE_ACTIVE
                    logging.info("I: received from %s:%d snapshot=%d" % (server.address, server.port, agent.sequence))
                else:
                    kvmsg.store(agent.kvmap)
            elif agent.state == STATE_ACTIVE:
                if kvmsg.sequence > agent.sequence:
                    agent.sequence = kvmsg.sequence
                    kvmsg.store(agent.kvmap)
                    action = "update" if kvmsg.body else "delete"
                    # logging.info("I: received from %s:%d %s=%d %s" % (server.address, server.port, action, agent.sequence, kvmsg.body))
        else:
            logging.info("I: server at %s:%d didn't give hugz" % (server.address, server.port))
            agent.cur_server = (agent.cur_server + 1) % len(agent.servers)
            agent.state = STATE_INITIAL